from os.path import join, abspath
import numpy as np
import random
import re
from typing import Set, List
from backend.experiment.framework.shots_prompt.task import \
    PerplexityPromptShotsChoiceMakingTaskMaker
from backend.experiment.framework.task import TaskMaker
from backend.experiment.framework.shots_prompt.data import ShotsGenerator
from backend.experiment.tokenizer import tokenizer
from backend.experiment.framework.shots_prompt.implement import \
    CommonPerplexityChoiceMakingSample
from backend.experiment.model import context_limit
from backend.experiment.clue.chid.scheme0.implement import CHIDImplement as S0


class CHIDStaticShotsGenerator(ShotsGenerator):
    def __init__(self, implement: 'CHIDImplement'):
        assert isinstance(implement, CHIDImplement)
        self.__implement = implement

    def get_shot(
            self, limit_remain: int, selected_shots: Set[int],
            **kwargs
    ) -> (np.ndarray, int) or None:
        shot_count = len(selected_shots)
        candidate = kwargs['candidate']
        if shot_count == 0:
            # 成语第一个示例返回成语的释义
            shot = self.__implement.idiom_meaning[candidate][1]
            assert len(shot) <= limit_remain, \
                f'{candidate}, {len(shot)}, {limit_remain}'
            return shot, -1
        else:
            assert candidate in self.__implement.shuffle_idiom_samples.keys()
            examples = self.__implement.shuffle_idiom_samples[candidate]
            next_select = len(selected_shots) - 1
            while next_select < len(examples):
                if next_select not in selected_shots:
                    if len(examples[next_select]) <= limit_remain:
                        return examples[next_select], next_select
                next_select += 1
            # 找不到长度比要求更小的了
            return None


class CHIDImplement(S0):
    __shuffle_idiom_samples = {}

    @property
    def shuffle_idiom_samples(self):
        if len(self.__shuffle_idiom_samples) == 0:
            for k, v in self.sorted_idiom_examples.items():
                new_order = v.copy()
                random.shuffle(new_order)
                self.__shuffle_idiom_samples[k] = new_order
        return self.__shuffle_idiom_samples

    def get_task_maker(self) -> TaskMaker:
        return PerplexityPromptShotsChoiceMakingTaskMaker(
            token_limit=context_limit,
            tokenizer=tokenizer,
            shots_generator=CHIDStaticShotsGenerator(implement=self)
        )

    @classmethod
    def _make_test_sample(cls, record, made_sample_num: int) -> \
            List[CommonPerplexityChoiceMakingSample]:
        res = []
        candidates: List[str] = record['candidates']
        contents: List[str] = record['content']
        for content in contents:
            found = re.findall('#idiom[0-9]*#', content)
            assert len(found) > 0
            if len(found) > 1:
                # 跳过有多个空的样例
                for _ in found:
                    res.append(None)
            else:
                prompt = content.replace(found[0], cls.__mask_token)
                res.append(
                    CommonPerplexityChoiceMakingSample(
                        candidates=candidates,
                        prompt=prompt,
                        sample_id=int(
                            found[0].replace('idiom', '').replace('#', '')),
                        target_mask_token=cls.__mask_token
                    )
                )
                made_sample_num += 1
        return res

    @classmethod
    def _finalize(cls):
        super()._finalize()
        cls.__shuffle_idiom_samples.clear()

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))
